from collections import UserString
from functools import cached_property, lru_cache
import math
from typing import ClassVar
from typing_extensions import Self


class Sequence(UserString):
    def reverse(self) -> Self:
        return self[::-1]

    def complement(self) -> Self:
        return self.translate(UserString.maketrans("atgcATGC", "tacgTACG"))

    def reverse_complement(self) -> Self:
        return self.reverse().complement()


class DonorSequence(Sequence):
    """splice donor, 9bp, 内含子3bp, 外显子6bp, 前3bp小写, 示例：cagGTAAGT"""

    Matrix: ClassVar[dict[str, float]]
    BGD: ClassVar[dict[str, float]] = {"A": 0.27, "C": 0.23, "G": 0.23, "T": 0.27}
    Seed1: ClassVar[dict[str, float]] = {"A": 0.004, "C": 0.0032, "G": 0.9896, "T": 0.0032}
    Seed2: ClassVar[dict[str, float]] = {"A": 0.0034, "C": 0.0039, "G": 0.0042, "T": 0.9884}

    @cached_property
    def formatted_sequence(self) -> Self:
        return self[:3].lower() + self[3:].upper()

    @lru_cache(maxsize=None)
    def score(self, default: float = -1.0) -> float:
        if len(self.data) == 9:
            # for key elements
            seq = self.formatted_sequence
            key = seq[3:5].upper()
            score = self.Seed1[key[0]] * self.Seed2[key[1]] / (self.BGD[key[0]] * self.BGD[key[1]])
            # for rest elements
            rest = (seq[:3] + seq[5:]).upper()
            rest_score = self.Matrix[rest]
            # final score
            return math.log(score * rest_score, 2)
        return default


class AcceptorSequence(Sequence):
    """splice donor, 23bp, 内含子3bp, 外显子20bp, 经典剪切区域2bp(内含子内)大写, 示例：ttccaaacgaacttttgtAGgga"""

    Matrix: ClassVar[dict[int, dict[int, float]]]
    BGD: ClassVar[dict[str, float]] = {"A": 0.27, "C": 0.23, "G": 0.23, "T": 0.27}
    Seed1: ClassVar[dict[str, float]] = {"A": 0.9903, "C": 0.0032, "G": 0.0034, "T": 0.0030}
    Seed2: ClassVar[dict[str, float]] = {"A": 0.0027, "C": 0.0037, "G": 0.9905, "T": 0.0030}

    @cached_property
    def formatted_sequence(self) -> Self:
        return self[:18].lower() + self[18:20].upper() + self[20:].lower()

    @property
    def hashseq(self) -> int:
        length = len(self)
        return sum(map(lambda x: int(x[1]) * 4 ** (length - x[0] - 1), enumerate(self.translate(self.maketrans("ACGT", "0123")))))

    @lru_cache(maxsize=None)
    def score(self, default: float = -1.0) -> float:
        if len(self.data) == 23:
            # for key elements
            seq = self.formatted_sequence
            key = seq[18:20].upper()
            score = self.Seed1[key[0]] * self.Seed2[key[1]] / (self.BGD[key[0]] * self.BGD[key[1]])
            # for rest elements
            rest = (seq[:18] + seq[20:]).upper()
            rest_score = 1
            rest_score *= self.Matrix[0][rest[:7].hashseq]
            rest_score *= self.Matrix[1][rest[7:14].hashseq]
            rest_score *= self.Matrix[2][rest[14:].hashseq]
            rest_score *= self.Matrix[3][rest[4:11].hashseq]
            rest_score *= self.Matrix[4][rest[11:18].hashseq]
            rest_score /= self.Matrix[5][rest[4:7].hashseq]
            rest_score /= self.Matrix[6][rest[7:11].hashseq]
            rest_score /= self.Matrix[7][rest[11:14].hashseq]
            rest_score /= self.Matrix[8][rest[14:18].hashseq]
            # final score
            return math.log(score * rest_score, 2)
        return default
